#pragma once
#include <atomic>
#include <omp.h>
#include <time.h>
#include "base.h"

namespace occ {
	class ScanlineZBuffer_Parallel {
	public:
		ScanlineZBuffer_Parallel(ScanlineZBufferDescription* scanline_desc) :
			bg(DefaultColor::Black()),
			framebuffer(nullptr),
			pMesh(nullptr),
			pLocal2world(nullptr),
			pWorld2view(nullptr),
			light_color(DefaultColor::White())
		{
			h = scanline_desc->h;
			w = scanline_desc->w;
			CPT.resize(h);
			zbuffer.resize(w * h);
			//zlock_line = new std::atomic_flag[w];

			perspective = getPerspectiveMatrix(
				scanline_desc->camera_desc->zNear,
				scanline_desc->camera_desc->zFar,
				scanline_desc->camera_desc->fovy,
				w, h);

			inverseZ = glm::identity<mat4f>();
			inverseZ[2].z = -1;

			light_pos = scanline_desc->light_pos;  // world coordinates
			//light_pos.z = -light_pos.z;

		}

		~ScanlineZBuffer_Parallel() {
			//delete[] zlock_line;
		}
		
		void setLightPosition(const vec3f& pos) {
			light_pos = vec3f(pos.x, pos.y, pos.z);
		}

		void setMesh(TriMesh* pMesh) { this->pMesh = pMesh; }
		void setLocal2world(mat4f* pLocal2world) { this->pLocal2world = pLocal2world; }
		void setWorld2view(mat4f* pWorld2view) { this->pWorld2view = pWorld2view; }
		void setFramebuffer(Color* framebuffer) { this->framebuffer = framebuffer; }
		void setBackgroundColor(const Color& bg) { this->bg = bg; }
		void unbindMesh() { pMesh = nullptr; }
		void unbindLocal2world() { pLocal2world = nullptr; }
		void unbindWorld2view() { pWorld2view = nullptr; }
		void unbindFramebuffer() { framebuffer = nullptr; }

		// called by the user. the pipeline does not reset framebuffer automatically.
		void resetFramebuffer() {
			std::fill(framebuffer, framebuffer + h * w, bg); 
		}

		// called by the user. the pipeline does not reset zbuffer automatically.
		void resetZbuffer() {
			std::fill(zbuffer.begin(), zbuffer.end(), 0);  // set zbuffer to zFar side
		}

		void runPipeline() {
			assert(pMesh);
			assert(framebuffer);
			processVertex();
			processPixel();
			postprocess();
		}

	protected:

		void resetStatus() {
			/* clear previous data */
			tri_pvs.clear();
			pvs_idx.clear();
			fi.clear();
			valid_indices.clear();
			CET.clear();
			for (int i = 0; i < h; ++i) {
				CPT[i].clear();
			}
			for (auto& vec : AET) {
				vec.clear();
			}
			for (auto& vec : AET_idx) {
				vec.clear();
			}
			for (auto& vec : AET_flag) {
				vec.clear();
			}
		}
		void processVertex() {
			mat4f trans = perspective;
			if (pWorld2view) trans *= *pWorld2view;
			trans *= inverseZ;
			if (pLocal2world) trans *= *pLocal2world;

			/* pass 1: pt => NDC space => clipping => screen space */
			const occFloat wh = w * 0.5;
			const occFloat hh = h * 0.5;

			tri_pvs.resize(pMesh->nFaces());
			fi.resize(pMesh->nFaces());
			std::fill(fi.begin(), fi.end(), 0);
			#pragma omp parallel for
			for (int i = 0; i < pMesh->nFaces(); ++i) {
				auto& face = pMesh->face[i];
				vec4f p1(pMesh->vp[face.x], 1);
				vec4f p2(pMesh->vp[face.y], 1);
				vec4f p3(pMesh->vp[face.z], 1);

				p1 = toNDC(trans, p1);
				p2 = toNDC(trans, p2);
				p3 = toNDC(trans, p3);
				if (inNDC(p1) || inNDC(p2) || inNDC(p3)) {
					// transform to screen space
					p1.x = wh * p1.x + wh;
					p2.x = wh * p2.x + wh;
					p3.x = wh * p3.x + wh;
					p1.y = wh * p1.y + wh;
					p2.y = wh * p2.y + wh;
					p3.y = wh * p3.y + wh;

					// use 1-z for precision. see RealTimeRendering v4, chap. 4
					p1.z = 1 - (p1.z + 1.0) * 0.5;
					p2.z = 1 - (p2.z + 1.0) * 0.5;
					p3.z = 1 - (p3.z + 1.0) * 0.5;

					tri_pvs[i] = Triangle(p1, p2, p3);
					fi[i] = 1;
				}
			}
			
			for (int i = 0; i < fi.size(); ++i) {
				if (fi[i]) valid_indices.push_back(i);
			}

			CET.resize(valid_indices.size() * 3);

			// prepare atomic cas flags
			std::atomic_flag* flags = new std::atomic_flag[h];
			for (int i = 0; i < h; ++i) flags[i].clear();

			/* pass 2: construct classified polygon and edge tables */
			#pragma omp parallel for num_threads(NUM_CORES)
			for (int i = 0; i < valid_indices.size(); ++i) {
				int idx = valid_indices[i];
				int original_tri_idx = idx;
				auto& tri = tri_pvs[idx];
				/* construct classified polygon table */
				PolygonInfo polyinfo = { Plane(tri.v1,tri.v2,tri.v3),original_tri_idx, i };
				int ymax = round(max(tri.v1.y, tri.v2.y, tri.v3.y));
				int ymin = round(min(tri.v1.y, tri.v2.y, tri.v3.y));
				bool is_out = ymax >= h;  // part of the triangle is above the screen
				if (ymax >= 0 && ymax - ymin > 0 && abs(polyinfo.plane.c) >= EPS) {
					int _h = !is_out ? ymax : h - 1;
					while (flags[_h].test_and_set()) {}
					CPT[_h].push_back(polyinfo);
					flags[_h].clear();
				}

				/* construct classified edge tables */
				EdgeInfo e[3];
				int ei0 = 0, ei1 = 1, ei2 = 2;
				for (int j = 0; j < 3; ++j) {
					vec3f* pu = &tri._v[j], * pl = &tri._v[(j + 1) % 3];
					if (pu->y < pl->y) pu = &tri._v[(j + 1) % 3], pl = &tri._v[j];
					occFloat dx = pu->x - pl->x;
					int eymax = round(pu->y);
					int eymin = round(pl->y);
					int dy = eymax - eymin;
					dx = dy != 0 ? -dx / (pu->y - pl->y) : 0;
					if (eymax != ymax || dy == 0) {
						ei2 = j;
						ei0 = (j + 1) % 3;
						ei1 = (j + 2) % 3;
					}
					if (eymax >= h) {
						dy = h - 1 - eymin;
						occFloat alpha = (pu->y - h) / (pu->y - pl->y);
						e[j].upper_x = pu->x * (1.0 - alpha) + pl->x * alpha;
					}
					else e[j].upper_x = pu->x;
					e[j].dx = dx;
					e[j].dy = dy;
					e[j].id = original_tri_idx;
					e[j].ymax = eymax;
				}
				if (e[ei0].upper_x > e[ei1].upper_x) {
					int t = ei0;
					ei0 = ei1;
					ei1 = t;
				}
				// ensure sequence: 
				//top left edge, top right edge, bottom edge/ top horizontal edge
				CET[3 * i + 0] = e[ei0];
				CET[3 * i + 1] = e[ei1];
				CET[3 * i + 2] = e[ei2];
			}
			delete[] flags;
			//printf("num CET = %d\n", CET.size());
		}

		void processPixel() {
			//mat4f trans = inverseZ;
			mat4f trans = glm::identity<mat4f>();
			if (pLocal2world) trans *= *pLocal2world;
			//if (pWorld2view) trans *= *pWorld2view;
			//mat4f trans_normal = glm::transpose(glm::inverse(trans));
			mat4f trans_normal = trans;


			#pragma omp parallel num_threads(NUM_CORES) 
			{
				int block_size = ceil(h / omp_get_num_threads());
				int tid = omp_get_thread_num();
				int _h = (tid + 1) * block_size < h ? (tid + 1) * block_size : h;
				for (int i = _h - 1; i >= 0; --i) {
					/* add active pair sets when finding new classified polygons */
					if(i >= tid * block_size) {
						for (auto& polyinfo : CPT[i]) {
							size_t id = 3 * polyinfo.edge_id;
							auto& e0 = CET[id + 0];
							auto& e1 = CET[id + 1];
							ActiveEdgeInfo ae;
							ae.xl = e0.upper_x;
							ae.dxl = e0.dx;
							ae.dyl = e0.dy;
							ae.xr = e1.upper_x;
							ae.dxr = e1.dx;
							ae.dyr = e1.dy;
							occFloat a = polyinfo.plane.a;
							occFloat b = polyinfo.plane.b;
							occFloat c = polyinfo.plane.c;
							occFloat d = polyinfo.plane.d;
							ae.dzx = -a / c;
							ae.dzy = b / c;
							ae.zl = (-d - a * ae.xl - b * i) / c;
							ae.id = polyinfo.prim_id;
							ae.next_edge = CET[id + 2].dy != 0 ? id + 2 : -1;
							AET[tid].emplace_back(ae);
							AET_flag[tid].push_back(1);
						}
					}

					/* update AET idx */
					AET_idx[tid].clear();
					for (int k = 0; k < AET_flag[tid].size(); ++k) {
						if (AET_flag[tid][k]) AET_idx[tid].push_back(k);
					}
					if (i<tid*block_size&&AET_idx[tid].size() == 0) break;
					for (int k = 0; k < AET_idx[tid].size(); ++k) {
						int idx = AET_idx[tid][k];
						auto& edgepair = AET[tid][idx];
						occFloat zx = edgepair.zl;
						int xl = round(edgepair.xl);
						int xr = round(edgepair.xr);

						//printf("line: %d, ae: %d, xl = %d, xr = %d\n", i, idx, xl, xr);
						/* incremental depth update */
						if (xl < 0) {
							zx += (-xl) * edgepair.dzx;
							xl = 0;
						}
						for (int j = xl; j <= xr; ++j) {
							if (j >= w) {
								break;
							}
							//acquireLock(j);
							int pix = w * i + w - j;
							if (zx > zbuffer[pix]) {
								zbuffer[pix] = zx;
								vec4f center_world = trans * vec4f(pMesh->face_center[edgepair.id], 1);
								vec4f normal_world = glm::normalize(trans_normal * vec4f(pMesh->face_normal[edgepair.id], 0));
								vec4f _light_pos = vec4f(light_pos, 1);
								//if (pWorld2view) _light_pos = *pWorld2view * _light_pos;
								Color color;
								faceFragmentShader(center_world, normal_world, _light_pos, light_color, pMesh->color, color);
								//facePhongFragmentShader(center_world, normal_world, _light_pos, light_color, pMesh->color, color);
								framebuffer[pix] = color;
							}
							//releaseLock(j);

							zx += edgepair.dzx;
						}
						/* update active edge pair */
						edgepair.dyl--;
						edgepair.dyr--;
						if ((edgepair.dyl) <= 0 && (edgepair.dyr) <= 0) {
							AET_flag[tid][idx] = 0;  // lazy del
							continue;
						}
						edgepair.xl += edgepair.dxl;

						edgepair.xr += edgepair.dxr;
						if ((edgepair.dyl <= 0 || edgepair.dyr <= 0) && edgepair.next_edge >= 0) {
							auto& e = CET[edgepair.next_edge];
							edgepair.next_edge = -1;
							if (edgepair.dyl == 0) {
								edgepair.xl = e.upper_x;
								edgepair.dxl = e.dx;
								edgepair.dyl = e.dy;
							}
							if (edgepair.dyr == 0) {
								edgepair.xr = e.upper_x;
								edgepair.dxr = e.dx;
								edgepair.dyr = e.dy;
							}
						}
						edgepair.zl += edgepair.dxl * edgepair.dzx + edgepair.dzy;
					}
				}
				
			}
		}

		void postprocess() { resetStatus(); }

	protected:
		int h, w;
		TriMesh* pMesh;
		mat4f* pLocal2world;
		mat4f* pWorld2view;
		Color* framebuffer;
		Color bg;
		Color light_color;
		vec3f light_pos;

		mat4f perspective;
		mat4f inverseZ;
		std::vector<Triangle> tri_pvs;  // potential visible set in sreen space
		std::vector<int> pvs_idx;
		std::vector<int> fi;  // face idx of pvs
		std::vector<int> valid_indices;
		std::vector<std::vector<PolygonInfo>> CPT;  // classified polygon table
		std::vector<EdgeInfo> CET;                  // classified edge table
		std::array<std::vector<ActiveEdgeInfo>, NUM_CORES> AET;            // active edge table
		std::array < std::vector<int>, NUM_CORES> AET_flag;
		std::array < std::vector<int>, NUM_CORES> AET_idx;
		std::vector<occFloat> zbuffer;
	};
}